#include "process.h"
#include "layout.h"
#include "mm.h"
#include "protect.h"
#include <string.h>

#define PROC_FREE	0
#define PROC_RUNNING	1
#define PROC_READY	2
#define PROC_BLOCKED	3
#define PROC_ZOMBIE	4

#define PW_WAITALL	-1
#define PW_NOWAIT	0

struct process *current;

#define PROC_NR_IDLE 1

#define SIZE_GDT	256
static struct descriptor *gdt_g = (struct descriptor *)ADDR_GDT;
static void init_descriptor(struct descriptor *desc, u32_t base, u32_t limit, u16_t attr)
{
	desc->limit_low		= limit & 0xFFFF;
	desc->base_low		= base & 0xFFFF;
	desc->base_mid		= (base >> 16) & 0xFF;
	desc->attr1		= attr & 0xFF;
	desc->limit_high_attr2	= (limit >> 16) | (attr >> 8);
	desc->base_high		= base >> 24;
}

static void set_tss_ldt(struct process *p,
		    u32_t eip, u32_t eflags,
		    u32_t esp, u32_t ebp,
		    u32_t page_dir)
{
	int i;

	/* Set TSS and LDT in GDT. */
	for(i = 6; i < SIZE_GDT; i += 2) {
		if(gdt_g[i].attr1 == 0)
			break;
	}

	init_descriptor(&gdt_g[i], (u32_t)&p->tss, sizeof(struct tss)-1, DA_386TSS);
	init_descriptor(&gdt_g[i+1], (u32_t)&p->ldt, sizeof(struct ldt)-1, DA_LDT);

	p->tss_selector = i * 8;
	p->ldt_selector = (i + 1) * 8 + 3;

	/* Set descriptors in LDT. */
	init_descriptor(&p->ldt.zero, 0, 0, 0);
	init_descriptor(&p->ldt.cs, 0, 0xfffff, DA_CR | DA_32 | DA_LIMIT_4K| DA_DPL3);
	init_descriptor(&p->ldt.ds, 0, 0xfffff, DA_DRW | DA_32 | DA_LIMIT_4K | DA_DPL3);
	init_descriptor(&p->ldt.ss, 0, 0xfffff, DA_DRW | DA_32 | DA_LIMIT_4K | DA_DPL3);
	init_descriptor(&p->ldt.ss_kernel, 0, 0xfffff, DA_DRW | DA_LIMIT_4K | DA_32);

	/* Set TSS. */
	p->tss.ss0 = SELECTOR_USER_STACK_K;
	p->tss.esp0 = (u32_t)p + 4092; // PAGE_SIZE - 4;
	p->tss.eip = eip;
	p->tss.eflags = eflags;

	p->tss.edi = 0;
	p->tss.esi = 0;
	p->tss.esp = esp;
	p->tss.ebp = ebp;
	p->tss.ebx = 0;
	p->tss.edx = 0;
	p->tss.ecx = 0;
	p->tss.eax = 0; // reture value of the child process

	p->tss.es = SELECTOR_USER_DATA;
	p->tss.cs = SELECTOR_USER_CODE;
	p->tss.ss = SELECTOR_USER_STACK;
	p->tss.ds = SELECTOR_USER_DATA;
	p->tss.fs = SELECTOR_USER_DATA;
	p->tss.gs = SELECTOR_VIDEO;

	p->tss.ldt_selector = p->ldt_selector;
	p->tss.cr3_pdbr = page_dir;
}

static void clear_tss_ldt(struct process *p)
{
	gdt_g[p->tss_selector / 8].attr1 = 0;
	gdt_g[p->ldt_selector / 8].attr1 = 0;
}

/* Return process by pid. */
static struct process *get_proc(int pid)
{
	struct process *p = current;
	do {
		if(p->pid == pid)
			return p;
		p = p->next;
	} while(p != current);

	return NULL;
}

static int enable_schedular;
void inline enable_sched(void)
{
	enable_schedular = 1;
}
void inline disable_sched(void)
{
	enable_schedular = 0;
}

/* Process Schedular.
 *
 * There're 2 cases that calling this function:
 * 1. The current process runs out its time slice(called by clock_int_handle);
 * 2. The current process blocks(called by pause()) or exits(by do_exit()).
 *
 * All processes(ready, blocked, and zombie) are in one list. This function
 * chooses one ready process from this list. If no ready, then choose IDLE.
 */
#define TIME_SLICE	10
void schedular()
{
	static int sched_timer = TIME_SLICE;
	struct process *p;

	sched_timer--;

	if(enable_schedular == 0)
		return;

	/* No need switch. */
	if(sched_timer > 0 && current->status == PROC_READY)
		return;

	sched_timer = TIME_SLICE;

	/* Choose a ready process. */
	for(p = current->next; p != current; p = p->next) {
		if(p->status == PROC_READY)
			break;
	}
	/* If no ready, choose IDLE. */
	if(p->status != PROC_READY)
		p = get_proc(PROC_NR_IDLE);
	/* If the next is current, no need switch. */
	if(p == current)
		return;

	/* Switch to the next process. */
	current = p;
	switch_to(current->tss_selector);
}

/* Initilize process struct. */
static struct process *init_process(u32_t eip, u32_t eflags,
		u32_t esp, u32_t ebp, u32_t page_dir)
{
	static int last_pid = PROC_NR_IDLE;
	struct process *p;
	int i;

	p = (struct process *)new_page(0);
	p->pid = last_pid++;
	p->father = current->pid;
	p->status = PROC_READY;

	set_tss_ldt(p, eip, eflags, esp, ebp, page_dir);

	p->work_dir = 1;/* TODO*/
	p->wait_pid = PW_NOWAIT;
	for(i = 0; i < NUM_OPENFILE; i++) {
		p->openf[i].inode_pt = NULL;
	}

	/* Insert into process list. */
	disable_sched();
	if(current) {
		p->prev = current;
		p->next = current->next;
		current->next->prev = p;
		current->next = p;
	} else { /* First process.*/
		p->next = p;
		p->prev = p;
	}
	enable_sched();
	return p;
}

static int clear_process(struct process *p, int *retval)
{
	int pid;

	pid = p->pid;
	if(retval)
		*retval = p->tss.eax;

	/* Delete from process list. */
	p->next->prev = p->prev;
	p->prev->next = p->next;

	/* Delete the process-struct page. */
	delete_page((u32_t)p);

	return pid;
}

#define ADDR_USER_STACK		(0x2000000-4)
/* This function should be called only as system call!!! */
int do_fork(u32_t z1, u32_t z2, u32_t z3, u32_t z4, u32_t z5, u32_t ebp,
	u32_t eip, u32_t cs, u32_t eflags, u32_t esp, u32_t ss)
{
	u32_t page_dir;

	page_dir = page_table_fork();

	return init_process(eip, eflags, esp, ebp, page_dir)->pid;
}

/* Create the first process, and turn to run it. */
void first_proc(void *func)
{
	current = init_process((u32_t)func, 0x202, ADDR_USER_STACK, 0, 0x1000);
	current->father = 0;
	sti_a();
	enable_sched();
	switch_to(current->tss_selector);
}

static void wakeup_p(struct process *p)
{
	if(p == NULL)
		return;
	if(p->status == PROC_BLOCKED)
		p->status = PROC_READY;
}

void wake_up(int pid)
{
	wakeup_p(get_proc(pid));
}

void pause(void)
{
	current->status = PROC_BLOCKED;
	enable_sched();
	schedular();
}

int do_wait(int *retval)
{
	struct process *p = current;
	int flag = 0;

	current->wait_pid = PW_WAITALL;

	disable_sched();
	do {
		if(p->father == current->pid) {
			flag = 1;
			if(p->status == PROC_ZOMBIE) {
				enable_sched();
				goto out;
			}
		}
		p = p->next;
	} while(p != current);
	enable_sched();

	if(flag == 0) {// no child
		current->wait_pid = PW_NOWAIT;
		return -1;
	}

	pause();
	p = get_proc(current->wait_pid);
	if(p->status != PROC_ZOMBIE)
		return -1;

out:
	current->wait_pid = PW_NOWAIT;
	return clear_process(p, retval);
}

int waitpid(int pid, int *retval, int block)
{
	struct process *p = get_proc(pid);

	if(p == NULL || p->father != current->pid)
		return -1;

	disable_sched();
	if(p->status != PROC_ZOMBIE) {
		if(block == 1)
			return 0;

		current->wait_pid = pid;
		pause();
		if(p->status != PROC_ZOMBIE)
			return -1;
	}
	enable_sched();

	return clear_process(p, retval);
}

/*
 * stack:
 * +---+
 * |   |
 * |   |
 * |   |
 * |   | <--- argments
 * |   |
 * |   | <--- argv_new
 * |   | <--- &argv_new
 * |   | <--- argc
 * |   | <--- return address (ESP)
 * |   |
 */
static u32_t set_args(char **argv_old)
{
	int argc, i;
	int len[10], length = 0;
	char *p, *buf;
	char **argv_new;
	u32_t *esp;

	/* Arguments. */
	/* We must alloc a new buffer to save argments,
	 * in order to avoid covering old argment by new one. */
	buf = new_page(0);
	p = buf;
	for(i = 0; argv_old[i]; i++) {
		strcpy(p, argv_old[i]);
		len[i] = strlen(p) + 1;
		p += len[i];
		length += len[i];
	}
	argc = i;
	esp = (u32_t *)((ADDR_USER_STACK - length) & ~3); /* align */
	memcpy(esp, buf, length);
	delete_page(buf);

	/* Pointers of arguments. */
	p = (char *)esp;
	esp -= argc;
	argv_new = (char **)esp;
	for(i = 0; i < argc; i++) {
		argv_new[i] = p;
		p += len[i];
	}

	*(--esp) = (u32_t)argv_new;	/* argv */
	*(--esp) = argc;		/* argc */
	*(--esp) = 0x87654321;		/* ret address, not used */
	return (u32_t)esp;
}

/*
 * Kernel satck:
 *
 * |          |
 * |    ss    |
 * |   esp *  |
 * |  eflags  |
 * |    cs    |
 * |   eip *  + int 80h
 * |   ebp    |
 * |   edi    |
 * |   esi    |
 * |   edx    |
 * |   ecx    |
 * |   ebx    |
 * | ret addr + call eax  <-- sp_kernel
 * |          |
*/
int do_exec(u32_t path, char **argv)
{
	u32_t *sp_kernel;
	u32_t eip, esp;

	/* Load elf. */
	if(load_elf(path, &eip) != 0)
		return 1;

	/* Set arguments. */
	esp = set_args(argv);

	/* Free shared memory. */
	free_page_table(1);
	flush_tlb();

	/* Set EIP and ESP. */
	sp_kernel = (u32_t *)(&path - 1);
	*(sp_kernel + 7) = eip;
	*(sp_kernel + 10) = esp;

	return 0;
}

/* Pseudo exec, jump to a function, not a ELF file. */
int do_exec1(u32_t begin_addr, char **argv)
{
	u32_t *esp;

	/* Set begin address. */
	esp = (u32_t *)&begin_addr;
	*(esp + 6) = begin_addr;

	/* Change esp. */
	*(esp + 9) += 4;

	return 0;
}

void do_exit(void)
{
	int i, j;
	u32_t *page_dir, *page_table;
	struct process *p, *tmp;

	/* Close open files. */
	for(i = 0; i < NUM_OPENFILE; i++)
		if (current->openf[i].inode_nr)
			close(i);

	/* Free memory. */
	free_page_table(0);
	delete_page(current->tss.cr3_pdbr);

	/* Handle son processes. */
	p = current;
	do {
		if(p->father == current->pid) {
			if(p->status == PROC_ZOMBIE) {
				/* Clear. */
				tmp = p->next;
				clear_process(p, 0);
				p = tmp;
			} else {
				/* Make it orphan. */
				p->father = PROC_NR_IDLE;
				p = p->next;
			}
		} else {
			p = p->next;
		}
	} while(p != current);

	disable_sched();

	/* Clear GDT. */
	clear_tss_ldt(current);

	/* Wake up father. */
	if(current->father == PROC_NR_IDLE) {
		/* Orphan. Clear itself. */
		tmp = current;
		/* Set 'current' for schedular. Can't use 'next', because it
		 * maybe 'idle', which will bring down schedular(). */
		current = current->prev;
		clear_process(tmp, 0);
	} else {
		p = get_proc(current->father);
		current->status = PROC_ZOMBIE;
		if(p->wait_pid == current->pid || p->wait_pid == PW_WAITALL) {
			p->wait_pid = current->pid;
			wake_up(current->father);
		}
	}
	enable_sched();

	schedular();
}

int getpid(void)
{
	return current->pid;
}

int getppid(void)
{
	return current->father;
}

void wait_for(struct process **head)
{
	current->wait_next = *head;
	*head = current;
}
void wake_wait(struct process **head)
{
	struct process *p;
	for(p = *head; p != NULL; p = p->wait_next)
		wakeup_p(p);
	*head = NULL;
}

void init_proc(void)
{
	memset((u32_t)(gdt_g + 6), 0, (SIZE_GDT - 6) * 8);
}
